Megatron Core MoE Updates - 2025 H2

颜子杰, 陈楷文 | NVIDIA GPU加速计算专家团队 | Nov 07, 2025

目录

• 概览

• MCore MoE 开发分支 (Dev Branch

  • 开发分支功能亮点

• MCore MoE 路线图

  • 下一版本路线图 (MCore v0.16

• 性能分析

  • DeepSeek-V3 端到端性能 - Hopper
  • DeepSeek-V3 端到端性能 - Blackwell
  • 长上下文训练性能 - Hopper
  • 进行中和未来的工作重点

• MCore MoE 优化之旅

  • MoE 并行折叠 (MoE Parallel Folding
  • DeepSeek-V3 在 H100 上的内存优化
  • 内存优化：细粒度激活卸载 (Fine-grained Activation Offloading
  • DeepSeek-V3 在 H100 上的算子融合 (Fusions
  • DeepSeek-V3 on H100: DeepEP
  • DeepSeek-V3 on H100: Blockwise FP8
  • 流水线并行优化：灵活的非对称流水线并行（PP）与自定义流水线布局
  • 交叉流水线并行（Interleaved PP）与A2A（All-to-All）重叠
  • DeepSeek-V3 on Hopper: EP A2A重叠与灵活VPP
  • DeepSeek-V3 on Hopper: MTP独立部署（MTP Standalone）
  • MCore中的MoE CUDA图（MoE CUDA Graph）
  • DeepSeek-V3性能最佳实践：一份配置清单
  • MCore MoE模型库 (Model Zoo

• Megatron-FSDP

  • 完全分片数据并行（FSDP）
  • Megatron-FSDP分片结构
  • Megatron-FSDP缓冲区管理
  • 用户缓冲区注册（User-Buffer-Registration）NCCL通信
  • 计算与通信的重叠
  • Megatron-FSDP 分布式环境
  • Megatron-FSDP 检查点 (Checkpointing
  • Megatron-FSDP 基准测试
  • PyTorch Conference 2025 演讲

概览

该报告概述了 Megatron Core (MCore) 中混合专家模型 (MoE) 的更新，重点介绍了在 NVIDIA Hopper 和 Blackwell 架构上实现的 DeepSeek-V3 SOTA 训练性能。

核心性能指标：

• H100: 370 TFLOPS
• B200: 730 TFLOPS
• GB200: 970 TFLOPS

下图展示了 MCore DeepSeek-V3 在不同硬件平台上的性能演进历史，从2025年初的170 TFLOPS逐步提升至2025年9月的970 TFLOPS。

实现这些性能的关键优化技术包括：
- 并行化 (Parallelism): 精细折叠 (Fine-grained Folding)
- 内存 (Memory): 精细重计算 (Fine-grained Recomputations)、内存高效置换 (Mem-Efficient Permutation)
- 通信 (Communication): DeepEP, HybridEP, 1F1B A2A 重叠
- CPU 开销 (CPU Overhead): 算子融合 (Fusions), cudaGraph
- FP8 (MXFP8/DeepSeek Recipe): 应用 FP8 精度进行训练
- 长上下文 (Long Context): 扩展至 512K 上下文长度

MCore MoE 开发分支 (Dev Branch)

为了支持新兴架构并实现最佳性能，MCore MoE 的开发在一个专用的开发分支上进行，其特点如下：
- 快速迭代与审查 (FAST ITERATION & REVIEW)
- SOTA MoE 特性 (SOTA MOE FEATURES)
- 全测试覆盖 (FULL TEST COVERAGE)
- 与主分支定期同步 (PERIODIC SYNC WITH MAIN)
- GitHub-First 开发模式 (GITHUB-FIRST DEVELOPMENT)
- 对社区贡献开放 (OPEN FOR CONTRIBUTIONS)

开发分支地址：https://github.com/NVIDIA/Megatron-LM/tree/dev

开发分支功能亮点

该开发分支已经合并了多项前沿功能：
- Qwen3-Next 和 GDN (Gated Delta Net) 支持 (已合并)
- 流水线并行下的激活卸载 (Activation Offloading) (已合并)
- Muon 和逐层分布 (layer-wise distribution) (已合并)
- 针对 NVL72 的 HybridEP (已合并)
- FSDP + EP 支持 (已合并)
- 重构和优化 MoE 的 cuda_graph_scope (已合并)
- MuonClip 支持
- 开发分支的混合上下文并行 (Hybrid-ContextParallel) 特性
- 融合的线性层和交叉熵 (Fused Linear and Cross Entropy) 操作

MCore MoE 路线图

以下是 MCore MoE 正在进行和未来的工作规划，涵盖模型支持、核心功能、并行化、优化、精度支持、通信和开发者体验等多个方面。

下一版本路线图 (MCore v0.16)

性能与内存增强
- [ ] 支持将 MTP 层放入独立的流水线阶段
- [ ] 融合线性和交叉熵操作

高级功能
- [ ] 增强 cuda_graph_scope 以支持 MoE 和 Mamba
- 更细粒度的图范围，如 MoE 路由器和调度预处理
- 微创式实现

• [ ] MuonClip 支持（非拆分版本）
• [ ] 为更简单的实现添加上下文并行支持

通信优化
- [ ] HybridEP MNNVL 优化
- [ ] HybridEP for NVL8+IB

Bug 修复
- [x] 修复 DeepSeek 和 Qwen HF tokenizer 的兼容性问题

持续进行的长期特性
- 针对 DeepSeek-V3, Qwen-3 等 fine-grained MoE 的端到端性能优化
- 无同步和全迭代 cudaGraph MoE 训练
- 设备发起的 HybridEP 和 GroupedGEMM
- MoE ECHO Dispatcher
- 针对 Blackwell 性能的 CPU 开销优化
- MLA CP 2.0 - 增强 MLA CP 以支持更长的序列训练
- 动态上下文并行以支持不平衡的长序列训练
- 探索 CPU 卸载 - Hybrid CPU-GPU expert routing
- Megatron FSDP MoE 训练的性能优化

性能分析

DeepSeek-V3 端到端性能 - Hopper

下表展示了在 Hopper 架构上，通过逐步应用各项优化技术，DeepSeek-V3 的性能提升过程。基线性能为 147 TFLOPS，最终通过 MTP Standalone 优化达到了 359 TFLOPS。

DeepSeek-V3 端到端性能 - Blackwell

下表详细列出了在 Blackwell 架构上，从基线 494.46 TFLOPS 到最终 970.01 TFLOPS 的各项优化特性及其带来的性能增益。关键优化包括 CUDA Graphs、HybridEP 和 Fuse quantization to normalization 等。

长上下文训练性能 - Hopper

下表展示了在 Hopper 平台上进行长上下文训练的性能数据，涵盖了不同 GPU 数量、序列长度和并行策略下的 TFLOPS 表现。

进行中和未来的工作重点

• Qwen3-Next 和 DeepSeek-V3.2 模型支持
• GB200 端到端性能交付
• 无同步和全迭代的 cudaGraph MoE 训练
• 设备发起的 HybridEP 和 GroupedGEMM
• MoE ECHO Dispatcher
• 全迭代 cudaGraph
• 更多算子融合
• 长上下文训练 (>512K)
• HybridEP
• 全面的 cudaGraph 支持
• Moun 和其他新兴优化器
• 利用 NVLINK-C2C 进行卸载

MCore MoE 优化之旅

MoE 并行折叠 (MoE Parallel Folding)

这是一个针对目标版本 v0.10 的技术。
- 核心思想: 解耦 MoE 层与 Attention 层的并行映射。允许在 Attention 和 MoE 层中使用不同的并行策略。
- 概念来源: MoE 并行折叠和 DLsim 中的异构并行 (Heterogeneous (UVXY) Parallelism) 本质上是相同的。我们默认将 MoE 中的 TPxEPxDP 从 Dense 和 TPxEPxDP 中折叠出来，并且没有为 CEP/TCP/TEP/CEP 提供显式设置。
- 接口设计:
- Attention 层使用传统的 TP/CP/DP。
- MoE 层可以指定任意的 TP/EP/DP。

• 用例:

  • Attention: TP8CP8DP1

  • MoE: 可以使用任意并行映射，例如：

    • TP1EP64DP1 (之前的扩展 EP)
    • TP64EP1DP1 (之前的扩展 TP)
    • TP2EP16DP2
    • EP8TP8DP1 (带重排序)

  • 之前的扩展 TP 和扩展 EP 只是 EP 折叠的两种特例。

DeepSeek-V3 在 H100 上的内存优化

• 基线使用完全重计算。内存优化旨在避免高开销的完全重计算。

• 细粒度重计算 (Fine-Grained Recomputing):

  • 以低开销减少内存使用。
  • 使用 "mla_up_proj" + "layernorm" + "moe_act" 的重计算策略。
  • 这可以在仅增加 3% 端到端开销的情况下，节省 33% 的内存。

• 内存高效置换 (Memory-Efficient Permutation):

  • 通过改变概率乘法的顺序，无需为反向传播保存 unpermutation 的输入。
  • 这可以在可忽略的开销下节省 20% 的内存。

• 我们还使用 TP2 进一步减少内存使用。

内存优化：细粒度激活卸载 (Fine-grained Activation Offloading)

• 在模块/算子粒度上卸载激活。
• 与细粒度重计算协同工作，以尽可能减少总激活内存。
• 支持 PP>1 & VPP>1。

• 与最新的优化和模型结构兼容：

  • MXFP8/NVFP4
  • 流水线并行布局
  • 1F1B A2A 重叠

• CUDA Graph (尚不能捕获卸载的模块)

• 混合模型 (尚未发布)

下图展示了在 Grace-Blackwell Trays 上运行 DeepSeek-V3 的性能，通过卸载 moe_act，吞吐量基本保持不变，但最大内存占用显著降低。

DeepSeek-V3 在 H100 上的算子融合 (Fusions)

• Kernel Fusions:

  • CE Loss fusion

  • MLA RoPE Fusion (用于前向和后向传播)

    • 前向: 955us -> 155us
    • 后向: 990us -> 150us

  • 我们还有其他融合，如置换融合 (permute fusion) 和路由融合 (router fusion)，与 Shopee 合作开发，但未包含在此次消融实验中。

下图展示了 MLA RoPE 融合前后的性能对比，以及融合操作对整体性能的提升。

DeepSeek-V3 on H100: DeepEP

DeepEP的优势包括：

• 减少专家并行（EP）的通信开销。
• 将全局排列（global permutation）与EP通信融合。

在使用Alltoall Dispatcher时，EP通信大约占据了50%的单步时间。DeepEP可以将EP通信时间大致减半。

下图展示了Alltoall Dispatcher和采用DeepEP的Flex Dispatcher的执行时间线对比，后者显著减少了通信（蓝色/紫色部分）的耗时。

性能数据

	Step Time (s)	Per GPU TFLOPs
DeepEP	32.7	261 (+54)

---

DeepSeek-V3 on H100: Blockwise FP8

• 本节介绍DeepSeek的分块FP8（Blockwise FP8）方案及所有与FP8相关的优化。
• FP8训练可以加速GEMM（通用矩阵乘法）计算并减小激活张量的大小。
• FP8主权重（primary weight）和BF16优化器状态可以极大地减少内存使用。
• 节省的内存对于后续的优化至关重要。

性能数据

	Step Time (s)	Per GPU TFLOPs
FP8	30.8	276

---

流水线并行优化：灵活的非对称流水线并行（PP）与自定义流水线布局

pipeline_model_parallel_layout 是一个用于定义流水线并行分区的灵活API，这对于非均衡模型的负载均衡至关重要。例如，对于DeepSeek-V3（61个解码器层 + 1个MTP层），我们可以通过设置 PP8VPP4 将其分区，具体命令如下：
--pipeline-model-parallel-layout="Et|(tt)*30|mL"

下图展示了如何将模型层（Embedding, Decoder, Loss）映射到不同的流水线阶段（Stage）以及不同的虚拟流水线并行（VPP）和物理流水线并行（PP）排名上，以实现负载均衡。

---

交叉流水线并行（Interleaved PP）与A2A（All-to-All）重叠

通过交叉执行不同微批次（micro-batch）的前向和后向计算，可以有效地将All-to-All通信操作与计算过程重叠，从而隐藏通信延迟。

• 下图中的块大小是根据nsys报告中提取的芯片执行时间确定的。
• 此工作与XHS合作完成。

---

DeepSeek-V3 on Hopper: EP A2A重叠与灵活VPP

• 启用专家并行（EP）的All-to-All（A2A）重叠，以隐藏EP通信。
• 使用灵活的虚拟流水线并行（VPP）以启用VPP4，从而减少流水线气泡（PP bubbles）。

下图的时间线展示了计算与通信的重叠情况，显著提升了GPU利用率。

性能数据

	Step Time (s)	Per GPU TFLOPs
EP Overlap	24.5	348 (+72)

<small>*仅供技术讨论和参考，性能可能因产品组合不同而异。</small>

---

DeepSeek-V3 on Hopper: MTP独立部署（MTP Standalone）

• 启用MTP Standalone优化。
• 在PP8VPP4配置下，总共有32个虚拟阶段。
• MTP层包含了embedding、decoder和loss的计算。将MTP和loss放在一起会使最后一个阶段的计算负载过重。通过优化层的放置，可以使不同PP排名上的工作负载更加均衡。

下表对比了有无MTP Standalone优化时各阶段的层分配情况：

通过将MTP和loss分离到不同的虚拟阶段，实现了更好的负载均衡。

性能数据

	Step Time (s)	Per GPU TFLOPs
MTP Standalone	23.8	359 (+11)

• 在不强制平衡的真实数据下，性能为：单步时间24.8秒，吞吐量345 TFLOPs。

<small>*仅供技术讨论和参考，性能可能因产品组合不同而异。</small>

---

MCore中的MoE CUDA图（MoE CUDA Graph）

通过cuda_graph_scope可以对CUDA图生效的范围进行细粒度控制。
* 可控范围包括：[attn, mlp, moe, moe_router, moe_preprocess]

例如，设置--cuda-graph-scope attn mlp moe_router moe_preprocess可以捕获整个密集层（dense layer）以及MoE层中的attn+router+preprocess部分。

下图左侧是MoE层的流程图，标出了不同部分的scope；右侧是性能分析器（profiler）的截图，显示了被CUDA图捕获（Graphed）的区域。

---

DeepSeek-V3性能最佳实践：一份配置清单

以下是一组推荐的命令行参数配置，用于实现DeepSeek-V3的最佳性能。

参数	功能
--moe-token-dispatcher-type flex --moe-enable-deepep	启用DeepEP
--moe-router-dtype fp32 --moe-permute-fusion --moe-grouped-gemm --moe-router-fusion	启用MoE相关优化
--pipeline-model-parallel-layout "Et*3|(tt)*29m|L"	使用灵活的PP布局
--recompute-granularity selective --recompute-modules mla_up_proj mlp moe_act	启用细粒度的重计算
--no-rope-fusion: false	启用MLA Rope Fusion
--cross-entropy-loss-fusion --cross-entropy-fusion-impl native	启用交叉熵损失函数融合
--manual-gc --manual-gc-interval 10	启用手动垃圾回收以避免迭代内的意外GC
--use-precision-aware-optimizer --exp-avg-dtype bf16 --exp-avg-sq-dtype bf16	启用BF16优化器状态
--overlap-moe-expert-parallel-comm --delay-wgrad-compute	启用1F1B重叠
--fp8-recipe blockwise --fp8-format e4m3	启用DeepSeek FP8方案
--moe-router-padding-for-fp8	启用MoE Router Padding优化
--fp8-param-gather	启用FP8主权重
--external-cuda-graph --cuda-graph-scope attn mlp moe_router moe_preprocess	启用CudaGraph（可选）

---

MCore MoE模型库 (Model Zoo)

为了方便客户复现我们的性能，我们提供了一个MCore MoE模型库。

• 包含性能复现所需的一切。
• 为典型的MoE模型（如DeepSeek, Qwen, Mixtral）提供预定义的端到端训练脚本。
• 预定义的模型参数。
• 经过调优的并行化和性能参数。
• 提供Dockerfiles。
• 链接: https://github.com/yanring/Megatron-MoE-ModelZoo

---

Megatron-FSDP

完全分片数据并行（FSDP）

为什么选择FSDP？

• 与3D并行相比，使用更简单。
• 与流水线并行（PP）及其衍生技术更兼容。
• 性能表现优异。

下图展示了在Blackwell平台上训练LM3 405B模型时，不同并行策略的性能对比。Megatron-FSDP (M-FSDP) 取得了最高的2020 TFLOPS/GPU。

---

Megatron-FSDP分片结构

FSDP2 vs Megatron-FSDP

• Torch FSDP2 均匀地对参数进行分片。
• Megatron-FSDP 非均匀地对模块进行分片。
• 两者都利用可扩展的按模块/按组的通信。
• FSDP2需要将按参数的分片（per-parameter shards）置换到按模块的通信缓冲区中，这在LM3 405B训练中导致了约10%的性能开销。

下图对比了两种分片方式的结构差异。

---

Megatron-FSDP缓冲区管理

FSDP双缓冲区（Double-Buffer）

• 基线FSDP：重复创建集合通信缓冲区。

  • 当剩余内存较少时，需要进行内存分配/释放，这会降低性能并增加内存碎片。
  • 会产生CPU开销。

• Megatron-FSDP：使用持久化的通信缓冲区。

  • 创建持久的双缓冲区，并在集合通信操作之间重用。
  • 在Blackwell平台上训练LM3 405B时，带来了1.5%的速度提升。

下图展示了Megatron-FSDP中使用预分配双缓冲区的流程。

---

用户缓冲区注册（User-Buffer-Registration）NCCL通信

• 基线集合通信核心（collective kernels）的缺点：

  • 占用8-32个流式多处理器（SMs），在与计算重叠时会拖慢计算核心。
  • 在核心和用户内存之间双向复制数据会带来额外的延迟。

• 注册持久化NCCL用户缓冲区的好处：

  • 启用零拷贝（zero-copy），消除拷贝延迟。
  • 能够利用NVIDIA NVLink和InfiniBand的特性。
  • 多播（Multi-cast）仅需1-4个SM，从而加速计算核心。
  • SHARP将归约（reduction）操作从GPU SM卸载到网络交换机。

下图展示了该技术所利用的硬件架构，包括GPU、NVSwitch和IB-Switch。

---

计算与通信的重叠

调整并行配置以隐藏通信

• 在计算中隐藏通信是FSDP性能的关键。
• 扩展计算量（例如，通过增加微批次大小）会导致更高的内存使用，但可以实现更好的重叠。

下图对比了基线反向传播和完全隐藏通信的两种情况。通过将微批次大小加倍（2X micro-batch size），计算时间足以完全覆盖通信时间，从而消除了暴露的通信开销。

Megatron-FSDP 分布式环境

Megatron-FSDP 的并行性由 DeviceMesh 定义和管理。

• Megatron-FSDP 的分布式状态与5个并行维度进行交互：

  • DP-Outer (节点间/Inter-Node)
  • DP-Shard (节点内/Intra-Node)
  • CP (上下文并行/Context Parallel)
  • TP (张量并行/Tensor Parallel)
  • EP (专家并行/Expert Parallel)，通过一个辅助的 "EDTP" DeviceMesh 支持。

• Megatron-FSDP 解耦了非 EP 模块和 EP 模块的分布式状态：

  • 对非 EP 模块强制使用 DeviceMesh。
  • 对 EP 模块使用辅助的 Expert DeviceMesh。

• Megatron-FSDP 与基于 Megatron 的并行性兼容。

Megatron-FSDP 检查点 (Checkpointing)

支持 Torch 分布式检查点 (DCP)

• 重构了检查点实现，以使用基于 DTensor 的 torch 分布式检查点。
• 支持将 3D 并行检查点转换为 Megatron-FSDP 检查点。

在 DeepSeek Proxy 模型上进行检查点重载测试

下图显示，从 Megatron-FSDP 和 3D 并行模式的检查点重新加载后，训练曲线与基线完全一致，证明了检查点加载的正确性和无缝衔接。

Megatron-FSDP 基准测试

Blackwell 平台

下表展示了在 Blackwell 平台上对 DeepSeek-V3 模型的基准测试结果。

• 启用 Megatron-FSDP 的标志 (flags)：

  • --use-megatron-fsdp
  • --data-parallel-sharding-strategy optim_grads_optim
  • --no-gradient-accumulation-fusion
  • --use-distributed-optimizer
  • --calculate-per-token-loss
  • --init-model-with-meta-device
  • --ckpt-format fsdp_dtensor

• 更高性能的选项：

  • --grad-reduce-in-bf16
  • --fsdp-double-buffer
  • --use-nccl-ub

• 未来可能的优化方向：

  • HybridEP (混合专家并行)
  • Activation offloading (激活卸载)
  • ...

PyTorch Conference 2025 演讲

该工作已在 PyTorch Conference 2025 上展示。

• 观看录播： https://youtu.be/EoePOh9bTKs?si=94aytyvDKOs3YQwL